///rosbag rewrite, and reorganize data by header timestamp
///support sensor_msgs::IMU CompressedImage,livox_ros_driver::CustomMsg and sensor_msgs::Imu
///Create Time: 2022.6.12
///Update Time: 2022.11.10: IMU msg preprocess, change the coordinate system and change degree to rad for helmet data
///Update Time: 2023.04.16: provide the option for organizing date by the receive order when recording. You can comment(uncomment) the bag_out.write() function for using it.
///Author: Weitong Wu, Wuhan University
///Contact: wwtgeomatics@gmail.com

#include <livox_ros_driver/CustomMsg.h>
#include <ros/ros.h>
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <sensor_msgs/CompressedImage.h>
#include <sensor_msgs/Imu.h>

#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/Imu.h>
#include <boost/foreach.hpp>
#include <ros/time.h>

using namespace std;

std::string bag_in_path;
std::string bag_out_path;
int lidar_type = 1;

int main(int argc, char **argv) {
    ros::init(argc, argv, "rosbag_edit_whu");
    ros::NodeHandle nh;

    nh.param<std::string>("input_bag_path", bag_in_path, "");
    nh.param<std::string>("output_bag_path", bag_out_path, "");
    nh.param<int>("lidar_type", lidar_type, 1);

    rosbag::Bag bag_in, bag_out;
    bag_in.open(bag_in_path);
    bag_out.open(bag_out_path, rosbag::bagmode::Write);

    ROS_INFO_STREAM("In data processing......");

    for (auto &m : rosbag::View(bag_in)) {
        if (m.isType<sensor_msgs::Imu>()) {
            sensor_msgs::Imu imu_msg = *m.instantiate<sensor_msgs::Imu>();
            sensor_msgs::Imu imu_msg_tmp = *m.instantiate<sensor_msgs::Imu>();
            if (lidar_type == 1)
            {
                //for Avia
                imu_msg.angular_velocity.x = -imu_msg_tmp.angular_velocity.y / 180.0 * M_PI;
                imu_msg.angular_velocity.y = -imu_msg_tmp.angular_velocity.z / 180.0 * M_PI;
                imu_msg.angular_velocity.z = imu_msg_tmp.angular_velocity.x / 180.0 * M_PI;
                imu_msg.linear_acceleration.x = -imu_msg_tmp.linear_acceleration.y;
                imu_msg.linear_acceleration.y = -imu_msg_tmp.linear_acceleration.z;
                imu_msg.linear_acceleration.z = imu_msg_tmp.linear_acceleration.x;
            }else{
                //for Mid70
                imu_msg.angular_velocity.x = imu_msg_tmp.angular_velocity.y / 180.0 * M_PI;
                imu_msg.angular_velocity.y = imu_msg_tmp.angular_velocity.z / 180.0 * M_PI;
                imu_msg.angular_velocity.z = imu_msg_tmp.angular_velocity.x / 180.0 * M_PI;
                imu_msg.linear_acceleration.x = imu_msg_tmp.linear_acceleration.y;
                imu_msg.linear_acceleration.y = imu_msg_tmp.linear_acceleration.z;
                imu_msg.linear_acceleration.z = imu_msg_tmp.linear_acceleration.x;
            }
            //bag_out.write("/imu0", imu_msg.header.stamp, imu_msg);
	    bag_out.write("/imu0", m.getTime(), imu_msg);
        }
        if (m.isType<livox_ros_driver::CustomMsg>()) {
            auto lidar_msg = m.instantiate<livox_ros_driver::CustomMsg>();
            //bag_out.write("/livox/lidar", lidar_msg->header.stamp, lidar_msg);
	    bag_out.write("/livox/lidar", m.getTime(), lidar_msg);
        }
        if(m.isType<sensor_msgs::CompressedImage>()){
            auto image_msg = m.instantiate<sensor_msgs::CompressedImage>();
            //bag_out.write("/camera0/compressed", image_msg->header.stamp, image_msg);
            bag_out.write("/camera0/compressed", m.getTime(), image_msg);
        }

    }
    
    ROS_INFO_STREAM("Data processing completed ! ");

    bag_in.close();
    bag_out.close();

}
